"""
RAG查询引擎统一接口

提供简化的API和向后兼容性，同时整合新的模块化功能。
"""

from typing import List, Optional, Dict, Any
from llama_index.core import VectorStoreIndex
from llama_index.core.schema import NodeWithScore
from llama_index.core.query_engine import BaseQueryEngine
from llama_index.core.response import Response

# 导入新的模块化组件
from .engine import RAGQueryEngine as CoreRAGQueryEngine, QueryResult
from .strategies import (
    RetrievalStrategy, RetrievalConfig, RetrievalMode,
    RetrievalStrategyFactory, create_default_config
)
from .retrievers import CustomRetriever, MultiRetriever, AdaptiveRetriever


class RAGQueryEngine(CoreRAGQueryEngine):
    """
    RAG查询引擎 - 统一接口

    提供向后兼容的API，同时支持新的高级功能。
    保持了原有接口的简洁性，同时增加了扩展性。
    """

    def __init__(self, index: Optional[VectorStoreIndex] = None, mode: RetrievalMode = RetrievalMode.STANDARD):
        """
        初始化查询引擎

        Args:
            index: 向量索引实例，默认使用全局索引管理器的索引
            mode: 检索模式，默认为标准模式
        """
        # 创建策略
        if index:
            # 如果提供了索引，直接使用默认策略
            strategy = RetrievalStrategyFactory.create_strategy(RetrievalMode.STANDARD)
            super().__init__(index, strategy)
        else:
            # 使用延迟初始化
            self._mode = mode
            super().__init__(index=index)

    def _initialize_engine(self) -> None:
        """初始化查询引擎（重写以支持模式选择）"""
        # 使用索引管理器的索引（如果未提供索引）
        if self._index is None:
            from ..data_storage import index_manager
            self._index = index_manager.index

        if self._index is None:
            raise RuntimeError("索引未加载，请先创建或加载索引")

        # 如果还没有策略，创建默认策略
        if self._strategy is None:
            self._strategy = RetrievalStrategyFactory.create_strategy(
                getattr(self, '_mode', RetrievalMode.STANDARD)
            )

        # 确保LLM已配置
        from ..generation import llm_manager
        llm_manager.configure_llama_index()

        # 创建检索器
        self._retriever = CustomRetriever(self._index, self._strategy)

        # 创建查询引擎
        self._query_engine = self._create_query_engine()

        print("RAG查询引擎初始化完成")

    # 保持向后兼容的方法

    def query_with_sources(self, query_text: str) -> Dict[str, Any]:
        """
        执行查询并返回详细信息（向后兼容）

        Args:
            query_text: 查询文本

        Returns:
            Dict[str, Any]: 包含查询结果和源信息的字典
        """
        result = self.query_with_details(query_text)
        return {
            'query': result.query,
            'answer': result.answer,
            'sources': result.sources,
            'response_metadata': result.response_metadata,
            'performance': {
                'retrieval_time': result.retrieval_time,
                'generation_time': result.generation_time,
                'total_time': result.total_time
            }
        }

    def similarity_search(self, query_text: str, top_k: int = 5) -> List[NodeWithScore]:
        """
        执行相似性搜索（向后兼容）

        Args:
            query_text: 查询文本
            top_k: 返回的结果数量

        Returns:
            List[NodeWithScore]: 相似节点列表
        """
        return super().similarity_search(query_text, top_k)

    # 新增的高级功能方法

    def set_strategy(self, mode: RetrievalMode, **kwargs) -> None:
        """
        设置检索策略

        Args:
            mode: 检索模式
            **kwargs: 策略参数
        """
        strategy = RetrievalStrategyFactory.create_strategy(mode, **kwargs)
        self.update_strategy(strategy)

    def set_config(self, **config_kwargs) -> None:
        """
        设置检索配置

        Args:
            **config_kwargs: 配置参数
        """
        config = RetrievalConfig(**config_kwargs)
        self.update_config(config)

    def create_multi_retriever(self, modes: List[RetrievalMode], weights: Optional[List[float]] = None) -> MultiRetriever:
        """
        创建多检索器

        Args:
            modes: 检索模式列表
            weights: 权重列表

        Returns:
            MultiRetriever: 多检索器实例
        """
        if self._index is None:
            self._initialize_engine()

        retrievers = []
        for mode in modes:
            strategy = RetrievalStrategyFactory.create_strategy(mode)
            retrievers.append(CustomRetriever(self._index, strategy))

        return MultiRetriever(retrievers, weights)

    def create_adaptive_retriever(self) -> AdaptiveRetriever:
        """
        创建自适应检索器

        Returns:
            AdaptiveRetriever: 自适应检索器实例
        """
        if self._index is None:
            self._initialize_engine()

        return AdaptiveRetriever(self._index)

    def get_available_strategies(self) -> Dict[str, str]:
        """
        获取可用的检索策略

        Returns:
            Dict[str, str]: 策略信息字典
        """
        return RetrievalStrategyFactory.get_strategy_info()

    def benchmark_query(self, query_text: str, modes: Optional[List[RetrievalMode]] = None) -> Dict[str, Any]:
        """
        基准测试不同策略的查询效果

        Args:
            query_text: 测试查询
            modes: 要测试的模式列表

        Returns:
            Dict[str, Any]: 基准测试结果
        """
        if modes is None:
            modes = list(RetrievalMode)

        results = {}
        original_strategy = self._strategy

        for mode in modes:
            try:
                # 切换策略
                self.set_strategy(mode)

                # 执行查询
                start_time = self._get_current_time()
                result = self.query_with_details(query_text)
                end_time = self._get_current_time()

                results[mode.value] = {
                    'answer': result.answer,
                    'source_count': len(result.sources),
                    'total_time': end_time - start_time,
                    'retrieval_time': result.retrieval_time,
                    'generation_time': result.generation_time
                }

            except Exception as e:
                results[mode.value] = {
                    'error': str(e),
                    'total_time': 0,
                    'source_count': 0
                }

        # 恢复原始策略
        if original_strategy:
            self.update_strategy(original_strategy)

        return {
            'query': query_text,
            'results': results,
            'best_strategy': self._find_best_strategy(results)
        }

    def _get_current_time(self) -> float:
        """获取当前时间"""
        import time
        return time.time()

    def _find_best_strategy(self, results: Dict[str, Any]) -> str:
        """找出最佳策略"""
        best_strategy = ""
        best_score = -1

        for mode, result in results.items():
            if 'error' in result:
                continue

            # 简单的评分：时间越短，源数量越合适，分数越高
            source_count = result['source_count']
            if 3 <= source_count <= 8:  # 理想的源数量范围
                source_score = 1.0
            else:
                source_score = max(0, 1.0 - abs(source_count - 5) / 5)

            time_score = max(0, 1.0 - result['total_time'] / 10)  # 假设10秒为基准

            total_score = source_score * 0.6 + time_score * 0.4

            if total_score > best_score:
                best_score = total_score
                best_strategy = mode

        return best_strategy


# 全局查询引擎实例（向后兼容）
rag_engine = RAGQueryEngine()

# 便捷函数
def create_engine(mode: RetrievalMode = RetrievalMode.STANDARD, **kwargs) -> RAGQueryEngine:
    """
    创建查询引擎的便捷函数

    Args:
        mode: 检索模式
        **kwargs: 其他参数

    Returns:
        RAGQueryEngine: 查询引擎实例
    """
    return RAGQueryEngine(mode=mode, **kwargs)


def create_custom_engine(strategy: RetrievalStrategy, **kwargs) -> RAGQueryEngine:
    """
    创建自定义查询引擎的便捷函数

    Args:
        strategy: 自定义检索策略
        **kwargs: 其他参数

    Returns:
        RAGQueryEngine: 查询引擎实例
    """
    from ..data_storage import index_manager
    index = index_manager.index
    return RAGQueryEngine(index=index).__init__(index, strategy)